# -*- coding: utf-8 -*-

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import gzip
import qmnist

# 定义模型常量
IMAGE_SIZE = 28
NUM_CHANNELS = 1
PIXEL_DEPTH = 255
NUM_LABELS = 1  
BASE_DIR = '/Users/patrick/Workspace/git/aphf_score_recognize/'


class CNN(object):
    """
    定义神经网络模型
    """
    def __init__(self):
        model = models.Sequential()
        # 第1层卷积，卷积核大小为3*3，32个，28*28为待训练图片的大小
        model.add(layers.Conv2D(32, (3, 3), input_shape=(28, 28, 1)))
        model.add(layers.Activation('relu'))
        model.add(layers.MaxPooling2D((2, 2)))

        # 第2层卷积，卷积核大小为3*3，64个
        model.add(layers.ZeroPadding2D((1,1)))
        model.add(layers.Conv2D(64, (3, 3)))
        model.add(layers.Activation('relu'))
        model.add(layers.BatchNormalization(epsilon=1e-6,axis=-1))
        model.add(layers.MaxPooling2D((2, 2)))

        # 第3层卷积，卷积核大小为3*3，64个
        model.add(layers.ZeroPadding2D((1,1)))
        model.add(layers.Conv2D(64, (3, 3)))
        model.add(layers.Activation('relu'))
        model.add(layers.BatchNormalization(epsilon=1e-6,axis=-1))
        model.add(layers.MaxPooling2D((2, 2)))

        model.add(layers.Dropout(0.25))
        model.add(layers.Flatten())

        model.add(layers.Dense(3168))
        model.add(layers.Activation('relu'))
        # model.add(layers.Dropout(0.5))
        model.add(layers.Dense(10))
        model.add(layers.Activation('softmax'))

        # model.summary()

        self.model = model   


class DataSource(object):
    """
    定义神经网络数据源
    """
    def __init__(self):
        # mnist数据集存储的位置，如何不存在将自动下载
        data_path = BASE_DIR + 'MNIST_data/mnist.npz'
        # mnist数据只用于测试
        _, (test_images, test_labels) = datasets.mnist.load_data(path=data_path)
        test_images = test_images.reshape((10000, 28, 28, 1))
        
        # 模型训练使用qmnist
        train_images = self.extract_data(BASE_DIR + 'MNIST_data/xnist-images-idx3-ubyte.gz', 402953)
        train_images = train_images.reshape((402953, 28, 28, 1))
        train_labels = self.extract_labels(BASE_DIR + 'MNIST_data/xnist-labels-idx2-int.xz')

        # 像素值映射到 0 - 1 之间
        train_images, test_images = train_images / 255.0, test_images / 255.0

        self.train_images, self.train_labels = train_images, train_labels
        self.test_images, self.test_labels = test_images, test_labels

    def extract_data(self, filename, num_images):
        """从qmnist中提取图片数据
        :param filename: qmnist图片文件路径
        :param num_images: 图片数量
        """
        print('Extracting from ', filename)
        with gzip.open(filename) as bytestream:
            bytestream.read(16)
            buf = bytestream.read(IMAGE_SIZE * IMAGE_SIZE *
                                    num_images * NUM_CHANNELS)
            data = np.frombuffer(buf, dtype=np.uint8).astype(np.float32)
            data = data.reshape(num_images, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS)
        return data
    
    def extract_labels(self, filename):
        """从qumnist中提取Label
        :param filename: qmnist Label文件路径 
        """
        labels = qmnist.read_idx2_int(filename)
        labels = np.reshape(labels[:, :1], (-1))
        return labels  